"""
MergeCheckpoints节点的工具函数
"""
import json
from typing import List, Dict
import asyncio
from .call_llm import call_llm_async
from .utils import parse_llm_json_response

def create_deduplication_prompt(checkpoints: list) -> str:
    """
    创建智能去重的提示词
    
    Args:
        checkpoints: 评估项列表
        
    Returns:
        str: 格式化的提示词
    """
    checkpoints_json = json.dumps(checkpoints, ensure_ascii=False, indent=2)
    
    return f"""
请对以下安全评估项进行智能去重，识别语义相似或重复的评估项。

去重规则：
1. 如果两个评估项的category和checkpoint内容基本相同或语义相似，则视为重复
2. 对于重复项，保留信息最完整、表述最准确的版本
3. 轻微的表述差异不算重复，但实质内容相同的要去重
4. 所有项目都属于同一个module，重点关注category和checkpoint的去重

请返回去重后的评估项列表，格式如下：
```json
[
  {{
    "module": "模块名称",
    "category": "分类名称",
    "checkpoint": "评估细项描述"
  }}
]
```

原始评估项列表：
{checkpoints_json}
"""

async def ai_deduplication_async(checkpoints: List[Dict]) -> List[Dict]:
    """
    使用AI进行智能去重
    
    Args:
        checkpoints: 原始评估项列表
        
    Returns:
        List[Dict]: 去重后的评估项列表
    """
    if len(checkpoints) <= 1:
        return checkpoints
        
    try:
        # 调用AI进行去重
        prompt = create_deduplication_prompt(checkpoints)
        response = await call_llm_async(prompt)
        
        # 解析结果
        unique_checkpoints = parse_llm_json_response(
            response=response,
            expected_type=list,
            default_value=None
        )
        
        if unique_checkpoints is None:
            print("AI返回格式错误，返回原始数据")
            return checkpoints
        
        # 验证每个项目都有必需字段
        for item in unique_checkpoints:
            if not all(key in item for key in ["module", "category", "checkpoint"]):
                print("AI返回数据不完整，返回原始数据")
                return checkpoints
        
        return unique_checkpoints
        
    except Exception as e:
        print(f"AI去重失败: {e}，返回原始数据")
        return checkpoints

def group_checkpoints_by_module(checkpoints: List[Dict]) -> Dict[str, List[Dict]]:
    """
    按module属性对评估项进行分组
    
    Args:
        checkpoints: 评估项列表
        
    Returns:
        Dict[str, List[Dict]]: 按module分组的字典
    """
    groups = {}
    for checkpoint in checkpoints:
        module = checkpoint.get('module', '未知模块')
        if module not in groups:
            groups[module] = []
        groups[module].append(checkpoint)
    return groups

async def process_single_module_group_async(module: str, checkpoints: List[Dict], semaphore: asyncio.Semaphore) -> List[Dict]:
    """
    异步处理单个module组的去重
    
    Args:
        module: 模块名称
        checkpoints: 该模块的评估项列表
        semaphore: 并发控制信号量
        
    Returns:
        List[Dict]: 去重后的评估项列表
    """
    if len(checkpoints) <= 1:
        return checkpoints
    
    async with semaphore:  # 控制并发数
        print(f"正在处理模块 '{module}': {len(checkpoints)} 个评估项")
        
        # 调用异步AI去重
        unique_checkpoints = await ai_deduplication_async(checkpoints)
        
        print(f"模块 '{module}' 去重完成: {len(checkpoints)} -> {len(unique_checkpoints)} 个评估项")
        return unique_checkpoints

async def ai_deduplication_by_groups_async(checkpoints: List[Dict], max_concurrency: int = 5) -> List[Dict]:
    """
    按module分组异步并行进行AI去重
    
    Args:
        checkpoints: 原始评估项列表
        max_concurrency: 最大并发数
        
    Returns:
        List[Dict]: 去重后的评估项列表
    """
    if not checkpoints:
        return checkpoints
    
    # 按module分组
    module_groups = group_checkpoints_by_module(checkpoints)
    print(f"按module分组完成，共 {len(module_groups)} 个模块")
    
    # 创建信号量控制并发
    semaphore = asyncio.Semaphore(max_concurrency)
    
    # 创建异步任务
    tasks = []
    for module, group_checkpoints in module_groups.items():
        task = process_single_module_group_async(module, group_checkpoints, semaphore)
        tasks.append((module, task))
    
    # 并发执行所有任务
    all_unique_checkpoints = []
    results = await asyncio.gather(*[task for _, task in tasks], return_exceptions=True)
    
    # 处理结果
    for i, result in enumerate(results):
        module = tasks[i][0]
        if isinstance(result, Exception):
            print(f"模块 '{module}' 处理失败: {result}")
            # 失败时使用原始数据
            all_unique_checkpoints.extend(module_groups[module])
        else:
            all_unique_checkpoints.extend(result)
    
    return all_unique_checkpoints
